from datetime import datetime, timedelta
import random
from sqlalchemy import create_engine, inspect, MetaData, Table, Column, Float, String, DateTime, text, Integer, JSON
from sqlalchemy.orm import sessionmaker
import sys
import os

# 获取当前文件所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 将server.py所在目录添加到系统路径
sys.path.append(current_dir)

# 导入数据模型
from server import RobotStatus, Task, Alert, MotionControl

# 数据库连接配置
DATABASE_URL = 'mysql+pymysql://root:Leg123456!@8.137.14.132/motor_db'

def create_session():
    engine = create_engine(DATABASE_URL)
    Session = sessionmaker(bind=engine)
    return Session()

def check_and_update_table(engine):
    print("检查并更新表结构...")
    inspector = inspect(engine)
    
    # 检查robot_status表是否存在
    if 'robot_status' in inspector.get_table_names():
        # 获取现有列
        columns = [col['name'] for col in inspector.get_columns('robot_status')]
        
        # 检查是否需要添加新列
        required_columns = {
            'battery_voltage': 'FLOAT NOT NULL DEFAULT 24.0',
            'battery_current': 'FLOAT NOT NULL DEFAULT 0.0',
            'battery_temperature': 'FLOAT NOT NULL DEFAULT 25.0',
            'motor_temperatures': 'JSON NOT NULL',
            'motor_currents': 'JSON NOT NULL',
            'motor_positions': 'JSON NOT NULL',
            'motor_speeds': 'JSON NOT NULL',
            'motor_torques': 'JSON NOT NULL',
            'motor_modes': 'JSON NOT NULL',
            'imu_data': 'JSON NOT NULL',
            'temperature': 'FLOAT',
            'running_status': 'VARCHAR(50)',
            'gps_latitude': 'FLOAT',
            'gps_longitude': 'FLOAT',
            'cpu_usage': 'FLOAT',
            'memory_usage': 'FLOAT',
            'created_at': 'DATETIME',
            'status_data': 'JSON NOT NULL DEFAULT (CAST(\'{}\' AS JSON))'  # Add default empty JSON object
        }
        
        # 添加缺失的列
        for col_name, col_type in required_columns.items():
            if col_name not in columns:
                print(f"添加列: {col_name}")
                if col_name == 'created_at':
                    # 对于created_at列，添加默认值为当前时间
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type} DEFAULT CURRENT_TIMESTAMP"))
                        connection.commit()
                elif col_name in ['motor_temperatures', 'motor_currents', 'motor_positions', 'motor_speeds', 'motor_torques']:
                    # 对于电机相关的数值JSON列，添加默认值
                    default_value = '[0.0, 0.0, 0.0, 0.0]'
                    if col_name == 'motor_temperatures':
                        default_value = '[25.0, 25.0, 25.0, 25.0]'
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type} DEFAULT (CAST('{default_value}' AS JSON))"))
                        connection.commit()
                elif col_name == 'motor_modes':
                    # 对于电机模式JSON列，添加默认值
                    default_value = '["standby", "standby", "standby", "standby"]'
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type} DEFAULT (CAST('{default_value}' AS JSON))"))
                        connection.commit()
                elif col_name == 'imu_data':
                    # 对于IMU数据JSON列，添加默认值
                    default_value = '{"acc": [0.0, 0.0, 0.0], "gyro": [0.0, 0.0, 0.0]}'
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type} DEFAULT (CAST('{default_value}' AS JSON))"))
                        connection.commit()
                elif col_name == 'status_data':
                    # 对于status_data列，添加默认值为空JSON对象
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type}"))
                        connection.commit()
                elif col_name in ['battery_voltage', 'battery_current', 'battery_temperature']:
                    # 对于电池相关的列，添加默认值
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type}"))
                        connection.commit()
                else:
                    with engine.connect() as connection:
                        connection.execute(text(f"ALTER TABLE robot_status ADD COLUMN {col_name} {col_type}"))
                        connection.commit()
    else:
        # 如果表不存在，创建表
        metadata = MetaData()
        Table('robot_status', metadata,
            Column('id', Integer, primary_key=True),
            Column('battery_level', Float, nullable=False),
            Column('battery_voltage', Float, nullable=False, default=24.0),
            Column('battery_current', Float, nullable=False, default=0.0),
            Column('battery_temperature', Float, nullable=False, default=25.0),
            Column('motor_temperatures', JSON, nullable=False),
            Column('motor_currents', JSON, nullable=False),
            Column('motor_positions', JSON, nullable=False),
            Column('motor_speeds', JSON, nullable=False),
            Column('motor_torques', JSON, nullable=False),
            Column('motor_modes', JSON, nullable=False),
            Column('imu_data', JSON, nullable=False),
            Column('temperature', Float, nullable=False),
            Column('running_status', String(50), nullable=False),
            Column('gps_latitude', Float),
            Column('gps_longitude', Float),
            Column('cpu_usage', Float),
            Column('memory_usage', Float),
            Column('created_at', DateTime, default=datetime.utcnow),
            Column('status_data', JSON, nullable=False, default={})
        )
        metadata.create_all(engine)
        print("创建robot_status表")

def generate_robot_status(session, count=100):
    print("生成机器狗状态数据...")
    base_time = datetime.now() - timedelta(days=7)
    
    for i in range(count):
        time_offset = timedelta(minutes=i*10)
        battery_level = random.uniform(20.0, 100.0)
        # 根据电池电量计算电压（假设满电为25V）
        battery_voltage = (battery_level / 100.0) * 25.0
        # 根据运行状态生成合理的电流值
        status_type = random.choice(['standby', 'walking', 'running', 'charging'])
        if status_type == 'standby':
            battery_current = random.uniform(0.1, 0.5)  # 待机电流小
            motor_currents = [random.uniform(0.0, 0.1) for _ in range(4)]  # 待机时电机电流很小
            motor_positions = [0.0 for _ in range(4)]  # 待机时电机位置为0
            motor_speeds = [0.0 for _ in range(4)]  # 待机时电机速度为0
            motor_torques = [0.0 for _ in range(4)]  # 待机时电机扭矩为0
            motor_modes = ["standby" for _ in range(4)]  # 待机模式
        elif status_type == 'walking':
            battery_current = random.uniform(1.0, 3.0)  # 行走电流中等
            motor_currents = [random.uniform(0.5, 1.5) for _ in range(4)]  # 行走时电机电流中等
            motor_positions = [random.uniform(-45.0, 45.0) for _ in range(4)]  # 行走时电机位置在±45度
            motor_speeds = [random.uniform(1.0, 3.0) for _ in range(4)]  # 行走时电机速度中等
            motor_torques = [random.uniform(0.5, 1.5) for _ in range(4)]  # 行走时电机扭矩中等
            motor_modes = ["walking" for _ in range(4)]  # 行走模式
        elif status_type == 'running':
            battery_current = random.uniform(2.5, 5.0)  # 跑步电流较大
            motor_currents = [random.uniform(1.5, 3.0) for _ in range(4)]  # 跑步时电机电流较大
            motor_positions = [random.uniform(-60.0, 60.0) for _ in range(4)]  # 跑步时电机位置在±60度
            motor_speeds = [random.uniform(3.0, 5.0) for _ in range(4)]  # 跑步时电机速度较大
            motor_torques = [random.uniform(1.5, 3.0) for _ in range(4)]  # 跑步时电机扭矩较大
            motor_modes = ["running" for _ in range(4)]  # 跑步模式
        else:  # charging
            battery_current = random.uniform(-3.0, -1.0)  # 充电电流为负值
            motor_currents = [0.0 for _ in range(4)]  # 充电时电机不工作
            motor_positions = [0.0 for _ in range(4)]  # 充电时电机位置为0
            motor_speeds = [0.0 for _ in range(4)]  # 充电时电机速度为0
            motor_torques = [0.0 for _ in range(4)]  # 充电时电机扭矩为0
            motor_modes = ["standby" for _ in range(4)]  # 待机模式
            
        # 根据运行状态和电流生成合理的温度值
        system_temp = random.uniform(20.0, 35.0)  # 基础系统温度
        if status_type == 'standby':
            battery_temp = system_temp + random.uniform(0, 2)  # 待机温度接近系统温度
            motor_temps = [system_temp + random.uniform(0, 1) for _ in range(4)]  # 四个电机温度接近系统温度
        elif status_type == 'charging':
            battery_temp = system_temp + random.uniform(3, 5)  # 充电时温度较高
            motor_temps = [system_temp + random.uniform(0, 2) for _ in range(4)]  # 充电时电机温度略高
        else:
            battery_temp = system_temp + abs(battery_current) * random.uniform(0.5, 1.0)  # 运动时温度随电流升高
            # 运动时电机温度升高更多，且可能有差异
            motor_temps = [
                system_temp + abs(curr) * random.uniform(1.0, 2.0) 
                for curr in motor_currents
            ]

        # 生成IMU数据
        if status_type == 'standby' or status_type == 'charging':
            imu_data = {
                "acc": [random.uniform(-0.1, 0.1) for _ in range(3)],  # 静止状态下的微小加速度
                "gyro": [random.uniform(-0.1, 0.1) for _ in range(3)]  # 静止状态下的微小角速度
            }
        else:
            imu_data = {
                "acc": [random.uniform(-2.0, 2.0) for _ in range(3)],  # 运动状态下的加速度
                "gyro": [random.uniform(-180.0, 180.0) for _ in range(3)]  # 运动状态下的角速度
            }
        
        status = RobotStatus(
            battery_level=battery_level,
            battery_voltage=battery_voltage,
            battery_current=battery_current,
            battery_temperature=battery_temp,
            motor_temperatures=motor_temps,  # 四个电机的温度值列表
            motor_currents=motor_currents,  # 四个电机的电流值列表
            motor_positions=motor_positions,  # 四个电机的位置值列表
            motor_speeds=motor_speeds,  # 四个电机的速度值列表
            motor_torques=motor_torques,  # 四个电机的扭矩值列表
            motor_modes=motor_modes,  # 四个电机的工作模式列表
            imu_data=imu_data,  # IMU数据
            temperature=system_temp,
            running_status=status_type,
            gps_latitude=39.9042 + random.uniform(-0.1, 0.1),
            gps_longitude=116.4074 + random.uniform(-0.1, 0.1),
            cpu_usage=random.uniform(10.0, 90.0),
            memory_usage=random.uniform(20.0, 80.0),
            created_at=base_time + time_offset,
            status_data={}  # 添加空的status_data字段
        )
        session.add(status)
    
    session.commit()
    print(f"成功生成 {count} 条机器狗状态数据")

def generate_tasks(session, count=50):
    print("生成任务数据...")
    task_types = ['patrol', 'monitor', 'test', 'maintenance']
    task_status = ['pending', 'running', 'completed', 'failed']
    base_time = datetime.now() - timedelta(days=7)
    
    for i in range(count):
        time_offset = timedelta(hours=i*2)
        status = random.choice(task_status)
        start_time = base_time + time_offset
        
        task = Task(
            name=f'Task_{i+1}',
            type=random.choice(task_types),
            status=status,
            start_time=start_time,
            end_time=start_time + timedelta(hours=1) if status in ['completed', 'failed'] else None,
            created_at=start_time
        )
        session.add(task)
    
    session.commit()
    print(f"成功生成 {count} 条任务数据")

def generate_alerts(session, count=30):
    print("生成告警数据...")
    alert_types = ['system', 'battery', 'temperature', 'motion', 'network']
    alert_levels = ['low', 'medium', 'high']
    alert_messages = [
        '电池电量低',
        '系统温度过高',
        '网络连接不稳定',
        '运动控制异常',
        '传感器数据异常',
        'CPU使用率过高',
        '内存使用率过高',
        'GPS信号丢失'
    ]
    base_time = datetime.now() - timedelta(days=7)
    
    for i in range(count):
        time_offset = timedelta(hours=i*4)
        alert = Alert(
            type=random.choice(alert_types),
            level=random.choice(alert_levels),
            message=random.choice(alert_messages),
            status=random.choice(['unprocessed', 'processed']),
            created_at=base_time + time_offset
        )
        session.add(alert)
    
    session.commit()
    print(f"成功生成 {count} 条告警数据")

def generate_motion_controls(session, count=80):
    print("生成运动控制数据...")
    gait_types = ['walk', 'run', 'climb', 'stop']
    directions = ['forward', 'backward', 'left', 'right', 'stop']
    base_time = datetime.now() - timedelta(days=7)
    
    for i in range(count):
        time_offset = timedelta(minutes=i*15)
        motion = MotionControl(
            gait_type=random.choice(gait_types),
            speed=random.uniform(0.0, 100.0),
            direction=random.choice(directions),
            created_at=base_time + time_offset
        )
        session.add(motion)
    
    session.commit()
    print(f"成功生成 {count} 条运动控制数据")

def main():
    try:
        engine = create_engine(DATABASE_URL)
        # 首先检查并更新表结构
        check_and_update_table(engine)
        
        session = create_session()
        
        # 生成各类测试数据
        generate_robot_status(session)
        generate_tasks(session)
        generate_alerts(session)
        generate_motion_controls(session)
        
        print("所有测试数据生成完成！")
        
    except Exception as e:
        print(f"生成测试数据时出错: {str(e)}")
    finally:
        if 'session' in locals():
            session.close()

if __name__ == '__main__':
    main() 